package cn.hserver.dns;

import cn.hserver.dns.bean.DnsBean;
import cn.hserver.dns.bean.DnsData;
import cn.hserver.dns.cache.DnsCache;
import cn.hserver.dns.cache.DnsDataCache;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.dns.*;

import java.util.List;

public class DnsHandler extends SimpleChannelInboundHandler<DatagramDnsQuery> {
    @Override
    public void channelRead0(ChannelHandlerContext ctx, DatagramDnsQuery query) {
        DatagramDnsResponse response = new DatagramDnsResponse(query.recipient(), query.sender(), query.id());
        try {
            DefaultDnsQuestion dnsQuestion = query.recordAt(DnsSection.QUESTION);
            response.addRecord(DnsSection.QUESTION, dnsQuestion);

            long l = System.currentTimeMillis();

            ByteBuf buf = null;
            List<DnsBean> search = DnsCache.getDomain(dnsQuestion.name());
            if (search == null) {
                search = DnsSearch.search(dnsQuestion.name());
                //缓存查询
                if (search != null) {
                    DnsCache.put(dnsQuestion.name(), search);
                }
            }


            if (search != null) {
                for (DnsBean dnsBean : search) {
                    if (dnsBean.getType() == dnsQuestion.type().intValue()) {
                        DnsDataCache.put(new DnsData(
                                dnsQuestion.name(),
                                dnsQuestion.type().name(),
                                dnsBean.getData(),
                                System.currentTimeMillis() - l
                        ));
                        String[] split = dnsBean.getData().split("\\.");
                        byte[] data = new byte[4];
                        for (int i = 0; i < split.length; i++) {
                            data[i] = (byte) Integer.parseInt(split[i]);
                        }
                        buf = Unpooled.wrappedBuffer(data);
                    }
                }
            }
            if (buf == null) {
                buf = Unpooled.wrappedBuffer(new byte[]{127, 0, 0, 1});
            }
            DefaultDnsRawRecord queryAnswer = new DefaultDnsRawRecord(dnsQuestion.name(),
                    DnsRecordType.A, 10, buf);
            response.addRecord(DnsSection.ANSWER, queryAnswer);

        } catch (Exception e) {
            System.out.println("查询域名异常：" + e);
        } finally {
            ctx.writeAndFlush(response);
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        cause.printStackTrace();
    }
}